use core::sync::atomic::{AtomicUsize, Ordering};

use super::{
    error::{KallocError, Result},
    page::{Page, __free_pages_memory},
    slab::slab_init,
    zone::{Zone, MAX_ORDER},
    GfpFlags,
};
use crate::{
    println,
    processor::{nr_cpus, this_processor_id},
    space::addr::{
        mem_setup, memblock_foreach, memstart, page_to_virt, pfn_to_phys, pfn_valid,
        phys_to_page_mut, phys_to_pfn, phys_to_virt, virt_to_page_mut, virt_to_pfn, PAGE_SHIFT,
    },
};

#[repr(C)]
pub(crate) enum ZoneType {
    Dma = 0,
    Normal,
    Max,
}

struct GlobalData {
    zones: [Zone; ZoneType::Max as usize],
    free_pages: AtomicUsize,
    reserved_pages: AtomicUsize,
    total_pages: AtomicUsize,
}

static GLOBAL_PAGE_ALLOC: GlobalData = GlobalData {
    zones: [Zone::new(), Zone::new()],
    free_pages: AtomicUsize::new(0),
    reserved_pages: AtomicUsize::new(0),
    total_pages: AtomicUsize::new(0),
};

pub(crate) fn zone_idx(ty: ZoneType) -> &'static Zone {
    &GLOBAL_PAGE_ALLOC.zones[ty as usize]
}

#[inline]
pub(crate) fn get_page_from_freelist(gfp_mask: GfpFlags, order: usize) -> Result<usize> {
    if order >= MAX_ORDER {
        return Err(KallocError::EORDER);
    }

    let ty = if gfp_mask.zone_normal() { ZoneType::Normal } else { ZoneType::Dma };
    let zone = zone_idx(ty);

    let mut page;
    let mut train = 0;
    loop {
        page = zone.rmqueue(order);

        if let Some(p) = page {
            unsafe {
                (*p).prep_new_page(order, gfp_mask);
            }
            return Ok(page_to_virt(unsafe { &*p }));
        }

        for cpu in 0..nr_cpus() {
            if cpu == this_processor_id() {
                continue;
            }
            train += zone.drain_pages_zone(cpu);
        }
        if train != 0 {
            continue;
        }
        return Err(KallocError::ENOMEM);
    }
}

#[inline]
pub(crate) fn free_pages_to_freelist(vaddr: usize, order: usize) {
    assert!(pfn_valid(virt_to_pfn(vaddr)));
    let page = virt_to_page_mut(vaddr);
    if order == 0 {
        page.free_page_to_pcp();
    } else {
        page.free_pages_to_area(order);
    }
}

#[inline(always)]
pub(crate) fn free_pages_sub(nr: usize) {
    GLOBAL_PAGE_ALLOC.free_pages.fetch_sub(nr, Ordering::Relaxed);
}

#[inline(always)]
pub(crate) fn free_pages_add(nr: usize) {
    GLOBAL_PAGE_ALLOC.free_pages.fetch_add(nr, Ordering::Relaxed);
}

fn free_pages_memory(start_pfn: usize, end_pfn: usize) {
    let offset = memstart() >> PAGE_SHIFT;
    let mut start = start_pfn + offset;
    let end = end_pfn + offset;
    while start < end {
        let mut order = (MAX_ORDER - 1).min(start.trailing_zeros() as usize);
        while start + (1 << order) > end {
            order -= 1;
        }

        let paddr = pfn_to_phys(start - offset);
        __free_pages_memory(phys_to_page_mut(paddr) as *mut Page, phys_to_virt(paddr), order);

        start += 1 << order;
    }
}

// 初始化内存分配器
pub(crate) fn kalloc_init() {
    mem_setup(|start, size, is_free| {
        let start_pfn = phys_to_pfn(start);
        let end_pfn = phys_to_pfn(start + size);
        for pfn in start_pfn..end_pfn {
            let page = phys_to_page_mut(pfn_to_phys(pfn));
            page.init_page(ZoneType::Normal, !is_free);
        }
        let nr = end_pfn - start_pfn;
        GLOBAL_PAGE_ALLOC.total_pages.fetch_add(nr, Ordering::Relaxed);
        if !is_free {
            GLOBAL_PAGE_ALLOC.reserved_pages.fetch_add(nr, Ordering::Relaxed);
        }
    });
    // dma 区目前没有分配任何可用内存, 全部释放到 normal 区
    let zone = &GLOBAL_PAGE_ALLOC.zones[ZoneType::Dma as usize];
    zone.init_zone(0);
    let zone = &GLOBAL_PAGE_ALLOC.zones[ZoneType::Normal as usize];
    zone.init_zone(nr_total_pages() - nr_reserved_pages());

    memblock_foreach(|start, size| {
        let start_pfn = phys_to_pfn(start);
        let end_pfn = phys_to_pfn(start + size);

        free_pages_memory(start_pfn, end_pfn);
    });

    slab_init();
}

/// 打印分配信息
pub fn kalloc_dump() {
    println!(
        "Mem total: {}M, reserved: {}M, free {}M",
        (nr_total_pages() << (PAGE_SHIFT - 10)) / 1024,
        (nr_reserved_pages() << (PAGE_SHIFT - 10)) / 1024,
        (nr_free_pages() << (PAGE_SHIFT - 10)) / 1024
    );
}

/// 总的空闲页内存数量
pub fn nr_free_pages() -> usize {
    GLOBAL_PAGE_ALLOC.free_pages.load(Ordering::Relaxed)
}

/// 总的预留内存页数量
pub fn nr_reserved_pages() -> usize {
    GLOBAL_PAGE_ALLOC.reserved_pages.load(Ordering::Relaxed)
}

/// 总的内存页数量
pub fn nr_total_pages() -> usize {
    GLOBAL_PAGE_ALLOC.total_pages.load(Ordering::Relaxed)
}

/// zone 区对应管理的内存页数量
pub fn zone_manage_pages(flags: GfpFlags) -> usize {
    let ty = if flags.zone_normal() { ZoneType::Normal } else { ZoneType::Dma };
    let zone = zone_idx(ty);
    zone.zone_managed_pages()
}
